import pymysql
import pandas as pd
import pymysql.cursors

class MysqlUtils(object):
    def __init__(self):
        # 修正：pymysql.connect（不是 cursors），创建数据库连接
        self.conn = pymysql.connect(
            host='127.0.0.1',
            user='root',
            password='root',
            database='tushare',
            port=3306,
            charset='utf8',
            cursorclass=pymysql.cursors.DictCursor  # 直接在连接时指定游标类型，后续无需重复设置
        )

class ClassificationUtils(object):
    def __init__(self):
        pass

    # 修正1：缩进为类方法，参数 seif → self
    def get_fina_indicator(self, conn):
        cursor = conn.cursor()  # 连接已指定 DictCursor，这里无需重复
        sql = """
        SELECT ts_code, ann_date, eps, total_revenue_ps, undist_profit_ps,
        gross_margin, fcff, fcfe, tangible_asset, bps, grossprofit_margin, npta
        FROM financial_data 
        WHERE ann_date BETWEEN '2023-01-01' and '2024-01-01'
        """
        cursor.execute(sql)
        rows = cursor.fetchall()
        df = pd.DataFrame(rows)
        
        # 剔除缺失关键字段的行
        df = df.dropna(subset=[
            'eps', 'total_revenue_ps', 'undist_profit_ps', 'gross_margin',
            'fcff', 'fcfe', 'tangible_asset', 'bps', 'grossprofit_margin', 'npta'
        ])
        df = df.reset_index(drop=True)
        
        cursor.close()  # 关闭游标，释放资源
        return df

    # 修正2：缩进为类方法，参数正确，修正变量名和字段名
    def get_daily(self, conn, df):
        cursor = conn.cursor()
        new_list = []
        
        for index, row in df.iterrows():
            print(f"正在处理第 {index} 条数据")
            ts_code = row['ts_code']
            # 修正：判断 ann_date 是否为 datetime 类型，避免 strftime 报错
            ann_date = row['ann_date'].strftime('%Y-%m-%d') if pd.notna(row['ann_date']) else ''
            if not ann_date:
                print(f"第 {index} 条数据缺少 ann_date，跳过")
                continue
            
            # 修正 SQL：统一大小写，规范别名使用
            sql = f"""
            SELECT trade_date, closes 
            FROM date_1 d 
            WHERE d.the_date > DATE('{ann_date}') 
              AND d.ts_code = '{ts_code}'
            ORDER BY d.trade_date ASC 
            LIMIT 20
            """
            
            try:
                cursor.execute(sql)
                rows = cursor.fetchall()
                df_daily = pd.DataFrame(rows)  # 变量名更清晰，避免与外部 df 冲突
                
                # 修正：df1 → df_daily，判断数据是否存在
                if len(df_daily) > 0:
                    max_closes = df_daily['closes'].max()
                    min_closes = df_daily['closes'].min()
                    the_closes = df_daily['closes'].iloc[-1]
                    
                    # 修正：字段名匹配（fcff/fcfe → 原查询字段，避免 KeyError）
                    new_list.append({
                        'ts_code': ts_code,
                        'ann_date': ann_date,
                        'max_closes': max_closes,
                        'min_closes': min_closes,
                        'the_closes': the_closes,
                        'eps': row['eps'],
                        'total_revenue_ps': row['total_revenue_ps'],
                        'undist_profit_ps': row['undist_profit_ps'],
                        'gross_margin': row['gross_margin'],
                        'fcff': row['fcff'],  # 原错误：fcf（数据库字段是 fcff）
                        'fcfe': row['fcfe'],  # 原错误：tcf（数据库字段是 fcfe）
                        'tangible_asset': row['tangible_asset'],
                        'bps': row['bps'],
                        'grossprofit_margin': row['grossprofit_margin'],
                        'npta': row['npta'],
                    })
            except Exception as e:
                print(f"第 {index} 条数据处理失败：{str(e)}")
                continue
        
        # 修正：循环结束后再写入 CSV，避免重复覆盖
        if new_list:
            df2 = pd.DataFrame(new_list)
            df2.to_csv('fina_indicator.csv', index=False, encoding='utf-8-sig')
            print(f"数据处理完成，共生成 {len(df2)} 条记录，已保存到 fina_indicator.csv")
        else:
            print("未获取到有效数据，未生成 CSV 文件")
        
        cursor.close()  # 关闭游标

if __name__ == '__main__':
    # 初始化工具类
    mu = MysqlUtils()
    classification_utils = ClassificationUtils()
    
    try:
        # 1. 获取财务指标数据
        df_fina = classification_utils.get_fina_indicator(mu.conn)
        print(f"成功获取 {len(df_fina)} 条财务指标数据")
        
        # 2. 获取对应股票的20日行情数据并生成CSV
        classification_utils.get_daily(mu.conn, df_fina)
    except Exception as e:
        print(f"程序执行失败：{str(e)}")
    finally:
        # 确保程序结束时关闭数据库连接
        mu.conn.close()
        print("数据库连接已关闭")
